!test

Review updated until commit be0e2ea

Description

• Implement RaggedIterDomain::combine() method to flatten ragged structures back to regular IterDomains

• Add Combine expression class to represent the inverse of Partition operation

• Include comprehensive validation for component-ragged pairing when Partition definition exists

• Provide extensive test coverage including basic combine, validation cases, and asNested integration

• Add detailed design document explaining validation strategy and architectural decisions

Val* combined_extent =
    IrBuilder::createInContainer<Val>(container, DataType::Index);

if (ragged->definition() != nullptr &&
    ragged->definition()->isA<Partition>()) {
  auto* partition = ragged->definition()->as<Partition>();
  IterDomain* expected_component = partition->component();

  NVF_ERROR(
      component == expected_component,
      "combine: component mismatch. The provided component does not match ",
      "the component from the Partition that created this "
      "RaggedIterDomain.\n",
      "  Provided component: ",
      component->toString(),
      "\n",
      "  Expected component: ",
      expected_component->toString());
}

NVF_ERROR_EQ(
    std::ranges::distance(extents_tv->getLogicalDomain() | TensorDomain::kNoReductions),
    1,
    "Unexpected rank of extent tensor: ",
    extents_tv->toString());

Greptile Summary

This PR introduces the Combine operation for RaggedIterDomain, which is the inverse of Partition. The operation combines a component IterDomain with a RaggedIterDomain back into a single regular IterDomain.

Key changes:

• Adds RaggedIterDomain::combine() static method that validates inputs and creates a combined IterDomain with symbolic extent
• Implements new Combine expression class that mirrors Partition structure (inputs: component + ragged, output: combined)
• Uses Option 3 validation approach: validates component-ragged pairing only when direct Partition definition exists, otherwise trusts user
• Unlike Partition which stores extents as an attribute, Combine accesses extents via ragged->extents() without storing it
• Creates symbolic extent Val for combined dimension rather than computing actual sum
• Includes comprehensive test coverage and detailed design documentation

The implementation follows nvFuser's design philosophy of trusting user-provided inputs (similar to arithmetic operations), with best-effort validation when feasible.

Confidence Score: 4/5

• This PR is safe to merge with low risk - well-designed foundation for ragged domain operations
• Score reflects solid implementation with thorough validation and testing, though symbolic extent approach defers some complexity to lowering phase. The Option 3 design choice is pragmatic and well-documented. No critical bugs found, but score is 4/5 because: (1) symbolic extent leaves relationship to actual sum implicit, which may complicate future lowering/indexing, and (2) inconsistency with Partition in not storing extents as attribute could cause issues if ragged domain loses connection to extents through transformations
• csrc/ir/internal_base_nodes.cpp - the symbolic extent approach should be monitored during lowering implementation to ensure proper index computation

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant User
    participant RaggedIterDomain
    participant Combine as Combine Expr
    participant IterDomain
    participant Partition as Partition Expr

    User->>RaggedIterDomain: combine(component, ragged)
    
    RaggedIterDomain->>RaggedIterDomain: Validate component != null
    RaggedIterDomain->>RaggedIterDomain: Validate ragged != null
    RaggedIterDomain->>RaggedIterDomain: Validate component is not RaggedIterDomain
    RaggedIterDomain->>RaggedIterDomain: Validate parallel types are Serial
    RaggedIterDomain->>RaggedIterDomain: Validate iter types are Iteration
    
    alt ragged has Partition definition
        RaggedIterDomain->>Partition: Get expected component
        Partition-->>RaggedIterDomain: Return component IterDomain
        RaggedIterDomain->>RaggedIterDomain: Validate component matches expected
    else No Partition definition
        Note over RaggedIterDomain: Trust user (Option 3)
    end
    
    RaggedIterDomain->>RaggedIterDomain: Get extents from ragged
    RaggedIterDomain->>RaggedIterDomain: Validate extents is 1D
    RaggedIterDomain->>RaggedIterDomain: Create symbolic extent Val
    
    RaggedIterDomain->>IterDomain: Create combined IterDomain
    IterDomain-->>RaggedIterDomain: Return combined_id
    
    RaggedIterDomain->>Combine: Create Combine expression
    Combine->>Combine: addOutput(combined_id)
    Combine->>Combine: addInput(component)
    Combine->>Combine: addInput(ragged)
    
    RaggedIterDomain-->>User: Return combined IterDomain

!test